import re
import cv2
import matplotlib.pyplot as plt
import numpy as np
import numpy as np
import argparse
import time
import cv2
import os
import cv2
import numpy as np
from numpy.core.numeric import zeros_like
from numpy.core.shape_base import block
import scipy.stats

def graph_cut(rgb_img, depth_img, mask, rect = (539+100, 60+100, 1036-100, 1020-100)):
    r = rect

    rgb_roi = rgb_img[int(r[1]):int(r[1] + r[3]), int(r[0]):int(r[0] + r[2])]
    depth_roi = depth_img[int(r[1]):int(r[1] + r[3]), int(r[0]):int(r[0] + r[2])]
    
    if np.all(mask[int(r[1]):int(r[1] + r[3]), int(r[0]):int(r[0] + r[2])] == 2):
        mask = np.zeros(rgb_roi.shape[:2], dtype=np.uint8)
        mask[:-1] = 2
        mask[0:20] = 0
        mask[-20:-1] = 0
        mask[:,0:20] = 0
        mask[:,-20:-1] = 0

        cv2.circle(mask, (int(mask.shape[0]/2), int(mask.shape[1]/2)), 6, 3, 12)
        cv2.circle(mask, (int(mask.shape[0]/2), int(mask.shape[1]/2)), 3, 1, 6)
    else:
        mask = mask[int(r[1]):int(r[1] + r[3]), int(r[0]):int(r[0] + r[2])]
        mask[0:20] = 0
        mask[-20:-1] = 0
        mask[:,0:20] = 0
        mask[:,-20:-1] = 0
    
    # 矩形roi
    rect = (int(r[0]), int(r[1]), int(r[2]), int(r[3]))  # 包括前景的矩形，格式为(x,y,w,h)
    bgdmodel = np.zeros((1, 65), np.float64)  # bg模型的临时数组
    fgdmodel = np.zeros((1, 65), np.float64)  # fg模型的临时数组
    r = (0, 0, rgb_roi.shape[0], rgb_roi.shape[1])
    cv2.grabCut(rgb_roi, mask, r, bgdmodel, fgdmodel, 5, mode=cv2.GC_INIT_WITH_MASK)
    
    # 提取前景和可能的前景区域
    mask2 = np.where((mask == 1) + (mask == 3), 1, 0).astype('uint8')
    mask_raw = mask2*255
    # cv2.cvtColor
    #轮廓
    contours,hierarchy = cv2.findContours(mask_raw,cv2.RETR_LIST,cv2.CHAIN_APPROX_SIMPLE)
    area = []
    for i in range(len(contours)):
        area.append(cv2.contourArea(contours[i]))
    max_idx = np.argmax(area)
    mask_raw = zeros_like(mask_raw)
    cv2.drawContours(mask_raw, contours, max_idx, 255, 0)
    cv2.fillPoly(mask_raw, [contours[max_idx]], 255)
    
    mask2 = mask_raw
    
    mask_depth = np.where((depth_roi < 900) * depth_roi, 1, 0).astype(np.uint8)
    mask2 = cv2.bitwise_and(mask2, mask_depth)

    rgb_result = cv2.bitwise_and(rgb_roi, rgb_roi, mask=mask2)
    rgb_result = cv2.bitwise_and(rgb_roi, rgb_roi, mask=mask2)
    depth_result = cv2.bitwise_and(depth_roi, depth_roi, mask=mask2)
    
    
    r = rect
    rgb_img[:] = 0
    rgb_img[int(r[1]):int(r[1] + r[3]), int(r[0]):int(r[0] + r[2])] = rgb_result
    depth_img[:] = 0
    depth_img[int(r[1]):int(r[1] + r[3]), int(r[0]):int(r[0] + r[2])] = depth_result
    
    return rgb_img, depth_img, mask2


def get_hist(img):
    # hist = [cv2.calcHist([img], [0], None, [256], [1, 255]),cv2.calcHist([img], [1], None, [255], [1, 255]),cv2.calcHist([img], [2], None, [255], [1, 255])]
    hist = [cv2.calcHist([img], [0], None, [256], [1, 255]),cv2.calcHist([img], [1], None, [256], [1, 255])]
    for j in range(len(hist)):
        hist[j] /= np.sum(hist[j])
    return hist
def get_histx(variety):
    image_name = f'{variety}.png'
    img_ori = cv2.imread(os.path.join('./template', image_name))
    img = cv2.cvtColor(img_ori, cv2.COLOR_BGR2HSV)
    histx = get_hist(img)
    return histx
def get_mask(img_ori, a_size, histx):
    rmask = (486+30, 914+30, 486+155+50-30, 914+167+100-30)
    rmask = (486-20, 914, 486+155+50+50+30, 914+167+100-30)
    r = rmask
    img = cv2.cvtColor(img_ori, cv2.COLOR_BGR2HSV)
    img = img[r[0]:r[2], r[1]:r[3]]
    # cv2.imshow('img', img)
    # cv2.waitKey(0)
    # img = img_ori
    mask_ori = np.ones(img_ori.shape[0:2], dtype=np.uint8)*2
    mask = np.ones(img.shape[0:2], dtype=np.uint8)*2
    for i in range(int(img.shape[0]/a_size)):
        for j in range(int(img.shape[1]/a_size)):
            r = (j*a_size, i*a_size, a_size, a_size)  #  (x_min, y_min, w, h)
            
            roi = img[int(r[1]):int(r[1] + r[3]), int(r[0]):int(r[0] + r[2])]
            histy = get_hist(roi)
            
            res = 0
            for k in range(len(histx)):
                res += scipy.stats.entropy(histy[k]+0.01, histx[k]+0.01)

            if res < 0.7 :
                xxx= 11111111111
                xxx= 2.1
                mark_a_size = 5
                down = (a_size - mark_a_size)/2
                roi = img[int(r[1] + down):int(r[1] + r[3] - down), int(r[0]+down):int(r[0] + r[2]-down)]
                histy = get_hist(roi)
                res = 0
                for k in range(len(histx)):
                    res += scipy.stats.entropy(histy[k]+0.001, histx[k]+0.001)
                if res < 3.7:
                    mask[int(r[1] + down):int(r[1] + r[3] - down), int(r[0]+down):int(r[0] + r[2]-down)] = 1
                    
    r = rmask
    mask_ori[r[0]:r[2], r[1]:r[3]] = mask
    return mask_ori

def mask_add_depth(img_depth, mask):
    mask_depth = np.where(img_depth < 880, 1, 0).astype(np.uint8)
    mask_raw = mask_depth*255
    contours,hierarchy = cv2.findContours(mask_raw,cv2.RETR_LIST,cv2.CHAIN_APPROX_SIMPLE)
    area = []
    for i in range(len(contours)):
        area.append(cv2.contourArea(contours[i]))
    max_idx = np.argmax(area)
    mask_raw = np.zeros_like(mask_raw)
    cv2.fillPoly(mask_raw, [contours[max_idx]], 255)
    # print(cv2.contourArea(contours[max_idx]))
    thick = cv2.contourArea(contours[max_idx])/20
    thick = max(thick, 10)
    thick = min(thick, 160)
    # print(thick)
    cv2.drawContours(mask_raw, contours, max_idx, 0, thick)
    mask_depth = mask_raw
    mask = np.where(mask_depth != 0, 1, mask)
    # cv2.imshow('depth', mask_depth*200)
    # cv2.waitKey(0)

    return mask

if __name__ == '__main__':
    x = 14
    img_rgb = cv2.imread(f'C:\\Users\\ys\\Desktop\\data_original\\RGBImages/RGB_{x}.png')
    img_depth = cv2.imread(f'C:\\Users\\ys\\Desktop\\data_original\\DepthImages/Depth_{x}.png', cv2.IMREAD_UNCHANGED)
    img_depth = np.where(img_depth == 0, 65535, img_depth)


    histx = get_histx('Satine')
    mask = get_mask(img_rgb, 15, histx)
    mask = mask_add_depth(img_depth, mask)
    mask1 = np.where(mask == 2, 0, 1).astype(np.uint8)
    mask2 = np.asarray([mask, mask, mask])
    mask2 = np.swapaxes(np.swapaxes(mask2, 1, 2), 0, 2)
    mask2 = np.where(mask2 == 2, 0, 1)
    # mask2 = np.where((mask == 2) + (mask == 0), 1, 0).astype('uint8')
    
    white_img = np.ones_like(img_rgb)*255

    result = cv2.bitwise_and(white_img, white_img, mask=mask1)
    result = np.swapaxes(np.swapaxes(result, 1, 2), 0, 1)
    result[1] *= 0
    result = np.swapaxes(np.swapaxes(result, 0, 1), 1, 2)
    rmask = (486-60, 914-60, 486+155+50+50, 914+167+100-30)
    r = rmask
    show = (img_rgb*0.5).astype(np.uint8)+(result*0.5).astype(np.uint8)
    
    show = show[r[0]:r[2], r[1]:r[3]]
    cv2.imshow('masked', show)
    cv2.imwrite('./result/show_full.jpg', show)

    rgb_img, depth_img, mask = graph_cut(img_rgb, img_depth, mask)
    cv2.imshow('img', rgb_img)
    cv2.imwrite('./result/result_full.jpg', rgb_img[r[0]:r[2], r[1]:r[3]])
    cv2.waitKey(0)
